import torch
from torch.utils.data import Dataset
from torchvision import transforms
import os
from PIL import Image
import numpy as np
import math

class TF_IDF_data(Dataset):
    def __init__(self, csv_path,idf_dict=None,is_train=True, idx=None):
        one_hot = False if idf_dict is None else True
        data_list = self.get_csv_data(csv_path,one_hot)

        rpt_word_dict = {}
        self.label_list = []
        self.report_id_list = []
        if idf_dict==None:
            self.idf_dict = {}
            for report_id, words, label in data_list:
                word_count = {}
                for word in words:
                    if words != '':
                        # word_set.add(word)
                        word_count[word] = word_count.get(word, 0) + 1
                for word in word_count:
                    self.idf_dict[word] = self.idf_dict.get(word, 0) + 1
                rpt_word_dict[report_id] = word_count
                self.label_list.append(label)
                self.report_id_list.append(report_id)
            for word in self.idf_dict:
                idf = self.idf_dict[word]
                idf = math.log(len(data_list) / idf)
                self.idf_dict[word] = idf
            idf_dict = self.idf_dict
        else:
            if not is_train:
                for report_id, words in data_list:
                    word_count = {}
                    for word in words:
                        if words != '':
                            word_count[word] = word_count.get(word, 0) + 1
                    rpt_word_dict[report_id] = word_count
                    self.report_id_list.append(report_id)
            else:
                for report_id, words, label in data_list:
                    word_count = {}
                    for word in words:
                        if words != '':
                            # word_set.add(word)
                            word_count[word] = word_count.get(word, 0) + 1
                    rpt_word_dict[report_id] = word_count
                    self.label_list.append(label)
                    self.report_id_list.append(report_id)

        self.tf_idf_list = np.array(self.get_tf_idf(idf_dict, rpt_word_dict))
        self.is_train = is_train
        if isinstance(idx, np.ndarray):
            self.tf_idf_list =self.tf_idf_list[idx,:]
    def __len__(self):
        return len(self.tf_idf_list)

    def __getitem__(self, index):
        report_id = self.report_id_list[index]
        if self.is_train:
            tf_idf = torch.Tensor(self.tf_idf_list[index])
            label = torch.Tensor(self.label_list[index])
            return report_id,tf_idf, label
        else:
            tf_idf = torch.Tensor(self.tf_idf_list[index])
            return report_id,tf_idf,
    def get_csv_data(self,csv_path,one_hot):
        data_list = []

        label_count={}
        with open(csv_path, 'r') as file:
            for line in file:
                groups = line.strip().split('|,|')
                report_id = groups[0]
                description = groups[1]
                words = description.split(' ')
                if len(groups)>=3:
                    label = groups[2]
                    areas = label.split(' ')
                    if one_hot:
                        label = np.zeros(17,dtype=np.int)

                        for area_id in areas:
                            if area_id!='':
                                label[int(area_id)] = 1
                                label_count[area_id] = label_count.get(area_id,0)+1
                    data_list.append((report_id, words, label))
                else:
                    data_list.append((report_id, words))
        print(sorted(label_count.items(),key=lambda x:x[0]))
        return data_list
    def get_tf_idf(self,idf_dict,rpt_word_dict):
        # print(rpt_word_dict)
        word_list = list(idf_dict.keys())
        rpt_tf_idf_list=[]
        for report_id in rpt_word_dict:
            word_count = rpt_word_dict[report_id]
            word_count_list = []
            for word in word_list:
                if word in word_count:
                    if word in idf_dict:
                        word_count_list.append(word_count[word]*idf_dict[word])
                    else:
                        word_count_list.append(0)
                else:
                    word_count_list.append(0)
            total_word = sum(word_count_list)
            tf_list = list(map(lambda x:x/total_word,word_count_list))
            rpt_tf_idf_list.append(tf_list)
        return rpt_tf_idf_list

if __name__ == "__main__":
    data = TF_IDF_data(r'../data/track1_round1_train_20210222.csv', is_train=True)
    print(data.__getitem__(1))
